





<!DOCTYPE html>
<html class="writer-html5" lang="zh-CN" >
<head>
  <meta charset="utf-8">
  
  <meta name="viewport" content="width=device-width, initial-scale=1.0">
  
  <title>编写自定义通行证 &mdash; tvm 0.8.dev1982 文档</title>
  

  
  <link rel="stylesheet" href="https://maxcdn.bootstrapcdn.com/bootstrap/4.0.0/css/bootstrap.min.css" integrity="sha384-Gn5384xqQ1aoWXA+058RXPxPg6fy4IWvTNh0E263XmFcJlSAwiGgFAW/dAiS6JXm" crossorigin="anonymous">
  <link rel="stylesheet" href="../../_static/css/theme.css" type="text/css" />
  <link rel="stylesheet" href="../../_static/pygments.css" type="text/css" />
  <link rel="stylesheet" href="../../_static/css/theme.css" type="text/css" />
  <link rel="stylesheet" href="../../_static/gallery.css" type="text/css" />
  <link rel="stylesheet" href="../../_static/pygments.css" type="text/css" />
  <link rel="stylesheet" href="../../_static/css/tlcpack_theme.css" type="text/css" />

  
  
    <link rel="shortcut icon" href="../../_static/tvm-logo-square.png"/>
  

  
  
  
  
    
      <script type="text/javascript" id="documentation_options" data-url_root="../../" src="../../_static/documentation_options.js"></script>
        <script data-url_root="../../" id="documentation_options" src="../../_static/documentation_options.js"></script>
        <script src="../../_static/jquery.js"></script>
        <script src="../../_static/underscore.js"></script>
        <script src="../../_static/doctools.js"></script>
        <script src="../../_static/translations.js"></script>
    
    <script type="text/javascript" src="../../_static/js/theme.js"></script>

    
    <script type="text/javascript" src="../../_static/js/tlcpack_theme.js"></script>
    <link rel="index" title="索引" href="../../genindex.html" />
    <link rel="search" title="搜索" href="../../search.html" />
    <link rel="next" title="How to Use TVM Pass Infra" href="use_pass_infra.html" />
    <link rel="prev" title="Extend TVM" href="index.html" /> 
</head>

<body class="wy-body-for-nav">

   
  <div class="wy-grid-for-nav">
    
    
<header class="header">
    <div class="innercontainer">
      <div class="headerInner d-flex justify-content-between align-items-center">
          <div class="headerLogo">
               <a href="https://tvm.apache.org/"><img src=https://tvm.apache.org/assets/images/logo.svg alt="logo"></a>
          </div>

          <div id="headMenu" class="headerNav">
            <button type="button" id="closeHeadMenu" class="navCloseBtn"><img src="../../_static/img/close-icon.svg" alt="Close"></button>
             <ul class="nav">
                <li class="nav-item">
                   <a class="nav-link" href=https://tvm.apache.org/community>Community</a>
                </li>
                <li class="nav-item">
                   <a class="nav-link" href=https://tvm.apache.org/download>Download</a>
                </li>
                <li class="nav-item">
                   <a class="nav-link" href=https://tvm.apache.org/vta>VTA</a>
                </li>
                <li class="nav-item">
                   <a class="nav-link" href=https://tvm.apache.org/blog>Blog</a>
                </li>
                <li class="nav-item">
                   <a class="nav-link" href=https://tvm.apache.org/docs>Docs</a>
                </li>
                <li class="nav-item">
                   <a class="nav-link" href=https://tvmconf.org>Conference</a>
                </li>
                <li class="nav-item">
                   <a class="nav-link" href=https://github.com/apache/tvm/>Github</a>
                </li>
                <li class="nav-item">
                   <a class="nav-link" href=https://tvmchinese.github.io/declaration_zh_CN.html>About-Translators</a>
                </li>
             </ul>
               <div class="responsivetlcdropdown">
                 <button type="button" class="btn-link">
                   ASF
                 </button>
                 <ul>
                     <li>
                       <a href=https://apache.org/>Apache Homepage</a>
                     </li>
                     <li>
                       <a href=https://www.apache.org/licenses/>License</a>
                     </li>
                     <li>
                       <a href=https://www.apache.org/foundation/sponsorship.html>Sponsorship</a>
                     </li>
                     <li>
                       <a href=https://www.apache.org/security/>Security</a>
                     </li>
                     <li>
                       <a href=https://www.apache.org/foundation/thanks.html>Thanks</a>
                     </li>
                     <li>
                       <a href=https://www.apache.org/events/current-event>Events</a>
                     </li>
                     <li>
                       <a href=https://www.zhihu.com/column/c_1429578595417563136>Zhihu</a>
                     </li>
                 </ul>
               </div>
          </div>
            <div class="responsiveMenuIcon">
              <button type="button" id="menuBtn" class="btn-menu"><img src="../../_static/img/menu-icon.svg" alt="Menu Icon"></button>
            </div>

            <div class="tlcDropdown">
              <div class="dropdown">
                <button type="button" class="btn-link dropdown-toggle" data-toggle="dropdown" aria-haspopup="true" aria-expanded="false">
                  ASF
                </button>
                <div class="dropdown-menu dropdown-menu-right">
                  <ul>
                     <li>
                       <a href=https://apache.org/>Apache Homepage</a>
                     </li>
                     <li>
                       <a href=https://www.apache.org/licenses/>License</a>
                     </li>
                     <li>
                       <a href=https://www.apache.org/foundation/sponsorship.html>Sponsorship</a>
                     </li>
                     <li>
                       <a href=https://www.apache.org/security/>Security</a>
                     </li>
                     <li>
                       <a href=https://www.apache.org/foundation/thanks.html>Thanks</a>
                     </li>
                     <li>
                       <a href=https://www.apache.org/events/current-event>Events</a>
                     </li>
                     <li>
                       <a href=https://www.zhihu.com/column/c_1429578595417563136>Zhihu</a>
                     </li>
                  </ul>
                </div>
              </div>
          </div>
       </div>
    </div>
 </header>
 
    <nav data-toggle="wy-nav-shift" class="wy-nav-side fixed">
      <div class="wy-side-scroll">
        <div class="wy-side-nav-search" >
          

          
            <a href="../../index.html">
          

          
            
            <img src="../../_static/tvm-logo-small.png" class="logo" alt="Logo"/>
          
          </a>

          
            
            
                <div class="version">
                  0.8.dev1982
                </div>
            
          

          
<div role="search">
  <form id="rtd-search-form" class="wy-form" action="../../search.html" method="get">
    <input type="text" name="q" placeholder="Search docs" />
    <input type="hidden" name="check_keywords" value="yes" />
    <input type="hidden" name="area" value="default" />
  </form>
</div>

          
        </div>

        
        <div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="main navigation">
          
            
            
              
            
            
              <p class="caption" role="heading"><span class="caption-text">如何开始</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../install/index.html">安装 TVM</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../contribute/index.html">贡献者指南</a></li>
</ul>
<p class="caption" role="heading"><span class="caption-text">用户引导</span></p>
<ul class="current">
<li class="toctree-l1"><a class="reference internal" href="../../tutorial/index.html">User Tutorial</a></li>
<li class="toctree-l1 current"><a class="reference internal" href="../index.html">How To Guides</a><ul class="current">
<li class="toctree-l2"><a class="reference internal" href="../compile_models/index.html">编译深度学习模型</a></li>
<li class="toctree-l2"><a class="reference internal" href="../deploy/index.html">TVM 部署模型和集成</a></li>
<li class="toctree-l2"><a class="reference internal" href="../work_with_relay/index.html">Work With Relay</a></li>
<li class="toctree-l2"><a class="reference internal" href="../work_with_schedules/index.html">Work With Tensor Expression and Schedules</a></li>
<li class="toctree-l2"><a class="reference internal" href="../optimize_operators/index.html">优化张量算子</a></li>
<li class="toctree-l2"><a class="reference internal" href="../tune_with_autotvm/index.html">Auto-Tune with Templates and AutoTVM</a></li>
<li class="toctree-l2"><a class="reference internal" href="../tune_with_autoscheduler/index.html">Use AutoScheduler for Template-Free Scheduling</a></li>
<li class="toctree-l2"><a class="reference internal" href="../work_with_microtvm/index.html">Work With microTVM</a></li>
<li class="toctree-l2 current"><a class="reference internal" href="index.html">Extend TVM</a><ul class="current">
<li class="toctree-l3 current"><a class="current reference internal" href="#">编写自定义通行证</a><ul>
<li class="toctree-l4"><a class="reference internal" href="#prerequisites">Prerequisites</a></li>
<li class="toctree-l4"><a class="reference internal" href="#writing-a-pass">Writing a Pass</a></li>
<li class="toctree-l4"><a class="reference internal" href="#glue-to-lowering">Glue to Lowering</a></li>
<li class="toctree-l4"><a class="reference internal" href="#quick-view">Quick View</a></li>
</ul>
</li>
<li class="toctree-l3"><a class="reference internal" href="use_pass_infra.html">How to Use TVM Pass Infra</a></li>
<li class="toctree-l3"><a class="reference internal" href="use_pass_instrument.html">How to Use TVM Pass Instrument</a></li>
<li class="toctree-l3"><a class="reference internal" href="bring_your_own_datatypes.html">Bring Your Own Datatypes to TVM</a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="../profile/index.html">Profile Models</a></li>
<li class="toctree-l2"><a class="reference internal" href="../../errors.html">Handle TVM Errors</a></li>
<li class="toctree-l2"><a class="reference internal" href="../../faq.html">常见提问</a></li>
</ul>
</li>
</ul>
<p class="caption" role="heading"><span class="caption-text">开发者引导</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../dev/tutorial/index.html">Developer Tutorial</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../dev/how_to/how_to.html">开发者指南</a></li>
</ul>
<p class="caption" role="heading"><span class="caption-text">架构指南</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../arch/index.html">Design and Architecture</a></li>
</ul>
<p class="caption" role="heading"><span class="caption-text">主题引导</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../topic/microtvm/index.html">microTVM：裸机使用TVM</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../topic/vta/index.html">VTA: Versatile Tensor Accelerator</a></li>
</ul>
<p class="caption" role="heading"><span class="caption-text">参考指南</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../reference/langref/index.html">语言参考</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../reference/api/python/index.html">Python API</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../reference/api/links.html">Other APIs</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../reference/publications.html">Publications</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../genindex.html">索引</a></li>
</ul>

            
          
        </div>
        
      </div>
    </nav>

    <section data-toggle="wy-nav-shift" class="wy-nav-content-wrap">
      
      <nav class="wy-nav-top" aria-label="top navigation" data-toggle="wy-nav-top">
        
            <div class="togglemenu">

            </div>
            <div class="nav-content">
              <!-- tvm -->
              Table of content
            </div>
        
      </nav>


      <div class="wy-nav-content">
        
        <div class="rst-content">
        

          




















<div role="navigation" aria-label="breadcrumbs navigation">

  <ul class="wy-breadcrumbs">
    
      <li><a href="../../index.html">Docs</a> <span class="br-arrow">></span></li>
        
          <li><a href="../index.html">How To Guides</a> <span class="br-arrow">></span></li>
        
          <li><a href="index.html">Extend TVM</a> <span class="br-arrow">></span></li>
        
      <li>编写自定义通行证</li>
    
    
      <li class="wy-breadcrumbs-aside">
        
            
            <a href="../../_sources/how_to/extend_tvm/low_level_custom_pass.rst.txt" rel="nofollow"> <img src="../../_static//img/source.svg" alt="viewsource"/></a>
          
        
      </li>
    
  </ul>

  
  <hr/>
</div>
          <div role="main" class="document" itemscope="itemscope" itemtype="http://schema.org/Article">
           <div itemprop="articleBody">
            
  <div class="sphx-glr-download-link-note admonition note">
<p class="admonition-title">注解</p>
<p>点击 <a class="reference internal" href="#sphx-glr-download-how-to-extend-tvm-low-level-custom-pass-py"><span class="std std-ref">这里</span></a> 下载全部示例代码</p>
</div>
<div class="sphx-glr-example-title section" id="writing-a-customized-pass">
<span id="sphx-glr-how-to-extend-tvm-low-level-custom-pass-py"></span><h1>编写自定义通行证<a class="headerlink" href="#writing-a-customized-pass" title="永久链接至标题">¶</a></h1>
<p><strong>作者</strong>: <a class="reference external" href="https://were.github.io">Jian Weng</a></p>
<p>TVM is a framework that abstracts away the heterogenity of machine learning accelerators.
Sometimes users may want customize some analysis and IR transformations
to adapt TVM to their own specialized hardware. This tutorial helps users write
a customized pass in TVM.</p>
<div class="section" id="prerequisites">
<h2>Prerequisites<a class="headerlink" href="#prerequisites" title="永久链接至标题">¶</a></h2>
<p>Before reading this tutorial, we assume readers have already known these topics well:</p>
<ul class="simple">
<li><p>Writing an algorithm in TVM and schedule it. Otherwise, see example tutorials like
<a class="reference internal" href="../optimize_operators/opt_gemm.html#opt-gemm"><span class="std std-ref">如何在CPU上优化GEMM（通用矩阵乘）</span></a>.</p></li>
<li><p>The basic structure of HalideIR. Otherwise, see <code class="docutils literal notranslate"><span class="pre">HalideIR/src/ir/IR.h</span></code> to learn what
attributes of IR nodes are defined.</p></li>
<li><p>Visitor design pattern. Otherwise, check the
<a class="reference external" href="https://docs.python.org/3/library/ast.html">Python AST module</a> to see how an AST
visitor is implemented.</p></li>
<li><p>How a Schedule is lowered to either an IRModule class or a LLVM module. Otherwise,
take a look at <code class="docutils literal notranslate"><span class="pre">python/tvm/build_module.py</span></code> to get some basics.</p></li>
</ul>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="kn">import</span> <span class="nn">tvm</span>
<span class="kn">from</span> <span class="nn">tvm</span> <span class="k">import</span> <span class="n">te</span>
<span class="kn">import</span> <span class="nn">numpy</span> <span class="k">as</span> <span class="nn">np</span>
</pre></div>
</div>
<p>We first write a very simple vector add and build it with the default schedule. Then, we use
our customized lowering pass to manipulate the IR directly instead of using schedule primitives.</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="n">n</span> <span class="o">=</span> <span class="n">tvm</span><span class="o">.</span><span class="n">tir</span><span class="o">.</span><span class="n">const</span><span class="p">(</span><span class="mi">128</span><span class="p">,</span> <span class="s2">&quot;int32&quot;</span><span class="p">)</span>
<span class="n">a</span> <span class="o">=</span> <span class="n">te</span><span class="o">.</span><span class="n">placeholder</span><span class="p">((</span><span class="n">n</span><span class="p">,),</span> <span class="n">name</span><span class="o">=</span><span class="s2">&quot;a&quot;</span><span class="p">)</span>
<span class="n">b</span> <span class="o">=</span> <span class="n">te</span><span class="o">.</span><span class="n">placeholder</span><span class="p">((</span><span class="n">n</span><span class="p">,),</span> <span class="n">name</span><span class="o">=</span><span class="s2">&quot;b&quot;</span><span class="p">)</span>
<span class="n">c</span> <span class="o">=</span> <span class="n">te</span><span class="o">.</span><span class="n">compute</span><span class="p">((</span><span class="n">n</span><span class="p">,),</span> <span class="k">lambda</span> <span class="n">i</span><span class="p">:</span> <span class="n">a</span><span class="p">[</span><span class="n">i</span><span class="p">]</span> <span class="o">+</span> <span class="n">b</span><span class="p">[</span><span class="n">i</span><span class="p">],</span> <span class="n">name</span><span class="o">=</span><span class="s2">&quot;c&quot;</span><span class="p">)</span>

<span class="n">sch</span> <span class="o">=</span> <span class="n">te</span><span class="o">.</span><span class="n">create_schedule</span><span class="p">(</span><span class="n">c</span><span class="o">.</span><span class="n">op</span><span class="p">)</span>
<span class="n">ir</span> <span class="o">=</span> <span class="n">tvm</span><span class="o">.</span><span class="n">lower</span><span class="p">(</span><span class="n">sch</span><span class="p">,</span> <span class="p">[</span><span class="n">a</span><span class="p">,</span> <span class="n">b</span><span class="p">,</span> <span class="n">c</span><span class="p">])</span>
<span class="nb">print</span><span class="p">(</span><span class="n">ir</span><span class="p">)</span>
</pre></div>
</div>
<p class="sphx-glr-script-out">输出:</p>
<div class="sphx-glr-script-out highlight-none notranslate"><div class="highlight"><pre><span></span>primfn(a_1: handle, b_1: handle, c_1: handle) -&gt; ()
  attr = {&quot;from_legacy_te_schedule&quot;: True, &quot;global_symbol&quot;: &quot;main&quot;, &quot;tir.noalias&quot;: True}
  buffers = {c: Buffer(c_2: Pointer(float32), float32, [128], []),
             a: Buffer(a_2: Pointer(float32), float32, [128], []),
             b: Buffer(b_2: Pointer(float32), float32, [128], [])}
  buffer_map = {a_1: a, b_1: b, c_1: c} {
  for (i: int32, 0, 128) {
    c_2[i] = ((float32*)a_2[i] + (float32*)b_2[i])
  }
}
</pre></div>
</div>
</div>
<div class="section" id="writing-a-pass">
<h2>Writing a Pass<a class="headerlink" href="#writing-a-pass" title="永久链接至标题">¶</a></h2>
<p>Essentially, an “IR transformation pass” is a function which maps a statement to a new statement.
Thus, we define this vectorize function and implement it step by step.</p>
<p>TVM already provides two class for users to both analyze and transform IR.</p>
<div class="section" id="ir-visitor">
<h3>IR Visitor<a class="headerlink" href="#ir-visitor" title="永久链接至标题">¶</a></h3>
<p>We can use <code class="docutils literal notranslate"><span class="pre">tvm.tir.stmt_functor.post_order_visit(stmt,</span> <span class="pre">func)</span></code> to gather information from the Halide IR.
<code class="docutils literal notranslate"><span class="pre">func</span></code> is a function callback. This function will be called before exiting the current IR node,
i.e. post-order visit. Then we leverage side effects to store the result of IR visit, because the
return value of <code class="docutils literal notranslate"><span class="pre">func</span></code> will be ignored.</p>
<div class="admonition note">
<p class="admonition-title">注解</p>
<p>You MUST use some array to store the result of IR visit. Even the value is a single variable.
This is mainly due to the constraints in the Python-C runtime. The variable values will be
refreshed every recursion but the array values will be preserved.</p>
</div>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="n">loops</span> <span class="o">=</span> <span class="p">[]</span>


<span class="k">def</span> <span class="nf">find_width8</span><span class="p">(</span><span class="n">op</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;Find all the &#39;tir.For&#39; nodes whose extent can be divided by 8.&quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">op</span><span class="p">,</span> <span class="n">tvm</span><span class="o">.</span><span class="n">tir</span><span class="o">.</span><span class="n">For</span><span class="p">):</span>
        <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">op</span><span class="o">.</span><span class="n">extent</span><span class="p">,</span> <span class="n">tvm</span><span class="o">.</span><span class="n">tir</span><span class="o">.</span><span class="n">IntImm</span><span class="p">):</span>
            <span class="k">if</span> <span class="n">op</span><span class="o">.</span><span class="n">extent</span><span class="o">.</span><span class="n">value</span> <span class="o">%</span> <span class="mi">8</span> <span class="o">==</span> <span class="mi">0</span><span class="p">:</span>
                <span class="n">loops</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">op</span><span class="p">)</span>
</pre></div>
</div>
</div>
<div class="section" id="ir-transformation">
<h3>IR Transformation<a class="headerlink" href="#ir-transformation" title="永久链接至标题">¶</a></h3>
<p>The transformation interface is slightly different from the visitor interface. There is only a
post-order callback in the visitor, but transformation visitor supports both a pre-order and a
post-order callback. If you want to keep the origin IR node, just return None. If you want to
change the current node to some node, use TVM IR maker interface to build it and return
this value.</p>
<div class="admonition note">
<p class="admonition-title">注解</p>
<p>If the pre-order function is called and returns a value which is not None, the post-order
function will be skipped.</p>
</div>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="k">def</span> <span class="nf">vectorize8</span><span class="p">(</span><span class="n">op</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;Split can vectorize the loops found in `find_width8`.&quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">op</span> <span class="ow">in</span> <span class="n">loops</span><span class="p">:</span>
        <span class="n">extent</span> <span class="o">=</span> <span class="n">op</span><span class="o">.</span><span class="n">extent</span><span class="o">.</span><span class="n">value</span>
        <span class="n">name</span> <span class="o">=</span> <span class="n">op</span><span class="o">.</span><span class="n">loop_var</span><span class="o">.</span><span class="n">name</span>
        <span class="n">lo</span><span class="p">,</span> <span class="n">li</span> <span class="o">=</span> <span class="n">te</span><span class="o">.</span><span class="n">var</span><span class="p">(</span><span class="n">name</span> <span class="o">+</span> <span class="s2">&quot;.outer&quot;</span><span class="p">),</span> <span class="n">te</span><span class="o">.</span><span class="n">var</span><span class="p">(</span><span class="n">name</span> <span class="o">+</span> <span class="s2">&quot;.inner&quot;</span><span class="p">)</span>
        <span class="n">body</span> <span class="o">=</span> <span class="n">tvm</span><span class="o">.</span><span class="n">tir</span><span class="o">.</span><span class="n">stmt_functor</span><span class="o">.</span><span class="n">substitute</span><span class="p">(</span><span class="n">op</span><span class="o">.</span><span class="n">body</span><span class="p">,</span> <span class="p">{</span><span class="n">op</span><span class="o">.</span><span class="n">loop_var</span><span class="p">:</span> <span class="n">lo</span> <span class="o">*</span> <span class="mi">8</span> <span class="o">+</span> <span class="n">li</span><span class="p">})</span>
        <span class="n">body</span> <span class="o">=</span> <span class="n">tvm</span><span class="o">.</span><span class="n">tir</span><span class="o">.</span><span class="n">For</span><span class="p">(</span><span class="n">li</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">8</span><span class="p">,</span> <span class="n">tvm</span><span class="o">.</span><span class="n">tir</span><span class="o">.</span><span class="n">ForKind</span><span class="o">.</span><span class="n">VECTORIZED</span><span class="p">,</span> <span class="n">body</span><span class="p">)</span>
        <span class="n">body</span> <span class="o">=</span> <span class="n">tvm</span><span class="o">.</span><span class="n">tir</span><span class="o">.</span><span class="n">For</span><span class="p">(</span><span class="n">lo</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="n">extent</span> <span class="o">//</span> <span class="mi">8</span><span class="p">,</span> <span class="n">tvm</span><span class="o">.</span><span class="n">tir</span><span class="o">.</span><span class="n">ForKind</span><span class="o">.</span><span class="n">SERIAL</span><span class="p">,</span> <span class="n">body</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">body</span>
    <span class="k">return</span> <span class="kc">None</span>


<span class="nd">@tvm</span><span class="o">.</span><span class="n">tir</span><span class="o">.</span><span class="n">transform</span><span class="o">.</span><span class="n">prim_func_pass</span><span class="p">(</span><span class="n">opt_level</span><span class="o">=</span><span class="mi">0</span><span class="p">)</span>
<span class="k">def</span> <span class="nf">vectorize</span><span class="p">(</span><span class="n">f</span><span class="p">,</span> <span class="n">mod</span><span class="p">,</span> <span class="n">ctx</span><span class="p">):</span>
    <span class="k">global</span> <span class="n">loops</span>

    <span class="n">tvm</span><span class="o">.</span><span class="n">tir</span><span class="o">.</span><span class="n">stmt_functor</span><span class="o">.</span><span class="n">post_order_visit</span><span class="p">(</span><span class="n">f</span><span class="o">.</span><span class="n">body</span><span class="p">,</span> <span class="n">find_width8</span><span class="p">)</span>

    <span class="k">if</span> <span class="ow">not</span> <span class="n">loops</span><span class="p">:</span>
        <span class="k">return</span> <span class="n">sf</span>

    <span class="c1"># The last list arugment indicates what kinds of nodes will be transformed.</span>
    <span class="c1"># Thus, in this case only `For` nodes will call `vectorize8`</span>
    <span class="k">return</span> <span class="n">f</span><span class="o">.</span><span class="n">with_body</span><span class="p">(</span><span class="n">tvm</span><span class="o">.</span><span class="n">tir</span><span class="o">.</span><span class="n">stmt_functor</span><span class="o">.</span><span class="n">ir_transform</span><span class="p">(</span><span class="n">f</span><span class="o">.</span><span class="n">body</span><span class="p">,</span> <span class="kc">None</span><span class="p">,</span> <span class="n">vectorize8</span><span class="p">,</span> <span class="p">[</span><span class="s2">&quot;tir.For&quot;</span><span class="p">]))</span>
</pre></div>
</div>
</div>
</div>
<div class="section" id="glue-to-lowering">
<h2>Glue to Lowering<a class="headerlink" href="#glue-to-lowering" title="永久链接至标题">¶</a></h2>
<p>So far, we are done with writing this IR transformation pass. What we need to do next is to glue
this pass to TVM’s lower pass.</p>
<p>In this case, we inject the pass written above into the TVM standard lowering
pass by feeding <strong>a list of tuple</strong> as argument to <code class="docutils literal notranslate"><span class="pre">tir.add_lower_pass</span></code>. “Tuple” indicates different
phases of lowering. In TVM, there are four phases of lowering and user-customized ones will be
called after each phase is done.</p>
<div class="admonition note">
<p class="admonition-title">注解</p>
<dl class="simple">
<dt>Here are the essential transformations done by each phase:</dt><dd><ul class="simple">
<li><p>Phase 0 generates the raw IR and loop levels.</p></li>
<li><p>Phase 1 flattens the array storage.</p></li>
<li><p>Phase 2 transforms loops, like unroll, vectorization and thread-binding.</p></li>
<li><p>Phase 3 does some cleanup work.</p></li>
</ul>
</dd>
</dl>
</div>
<p>Thus, a good place to put this transformation pass is just after Phase 1.</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="k">with</span> <span class="n">tvm</span><span class="o">.</span><span class="n">transform</span><span class="o">.</span><span class="n">PassContext</span><span class="p">(</span><span class="n">config</span><span class="o">=</span><span class="p">{</span><span class="s2">&quot;tir.add_lower_pass&quot;</span><span class="p">:</span> <span class="p">[(</span><span class="mi">1</span><span class="p">,</span> <span class="n">vectorize</span><span class="p">)]}):</span>
    <span class="nb">print</span><span class="p">(</span><span class="n">tvm</span><span class="o">.</span><span class="n">lower</span><span class="p">(</span><span class="n">sch</span><span class="p">,</span> <span class="p">[</span><span class="n">a</span><span class="p">,</span> <span class="n">b</span><span class="p">,</span> <span class="n">c</span><span class="p">]))</span>
</pre></div>
</div>
<p class="sphx-glr-script-out">输出:</p>
<div class="sphx-glr-script-out highlight-none notranslate"><div class="highlight"><pre><span></span>primfn(a_1: handle, b_1: handle, c_1: handle) -&gt; ()
  attr = {&quot;from_legacy_te_schedule&quot;: True, &quot;global_symbol&quot;: &quot;main&quot;, &quot;tir.noalias&quot;: True}
  buffers = {a: Buffer(a_2: Pointer(float32), float32, [128], []),
             c: Buffer(c_2: Pointer(float32), float32, [128], []),
             b: Buffer(b_2: Pointer(float32), float32, [128], [])}
  buffer_map = {a_1: a, b_1: b, c_1: c} {
  for (i.outer: int32, 0, 16) {
    c_2[ramp((i.outer*8), 1, 8)] = ((float32x8*)a_2[ramp((i.outer*8), 1, 8)] + (float32x8*)b_2[ramp((i.outer*8), 1, 8)])
  }
}
</pre></div>
</div>
</div>
<div class="section" id="quick-view">
<h2>Quick View<a class="headerlink" href="#quick-view" title="永久链接至标题">¶</a></h2>
<p>This tutorial gives a quick view of writing a customized IR transformation pass:
- Use <code class="docutils literal notranslate"><span class="pre">tvm.tir.stmt_functor.post_order_visit</span></code> to gather information on each IR nodes.
- Use <code class="docutils literal notranslate"><span class="pre">tvm.tir.stmt_functor.ir_transform</span></code> to transform IR nodes.
- Wrap up two above to write an IR-transformation function.
- Use <code class="docutils literal notranslate"><span class="pre">tvm.transform.PassContext</span></code> to put this function to TVM lowering pass</p>
<div class="sphx-glr-footer class sphx-glr-footer-example docutils container" id="sphx-glr-download-how-to-extend-tvm-low-level-custom-pass-py">
<div class="sphx-glr-download docutils container">
<p><a class="reference download internal" download="" href="../../_downloads/caa649473e845a115a0397a2855fd356/low_level_custom_pass.py"><code class="xref download docutils literal notranslate"><span class="pre">Python</span> <span class="pre">源码下载:</span> <span class="pre">low_level_custom_pass.py</span></code></a></p>
</div>
<div class="sphx-glr-download docutils container">
<p><a class="reference download internal" download="" href="../../_downloads/d58ec306b89044968adefb49e6552378/low_level_custom_pass.ipynb"><code class="xref download docutils literal notranslate"><span class="pre">Jupyter</span> <span class="pre">notebook</span> <span class="pre">下载:</span> <span class="pre">low_level_custom_pass.ipynb</span></code></a></p>
</div>
</div>
<p class="sphx-glr-signature"><a class="reference external" href="https://sphinx-gallery.github.io">Gallery generated by Sphinx-Gallery</a></p>
</div>
</div>


           </div>
           
          </div>
          

<footer>

    <div class="rst-footer-buttons" role="navigation" aria-label="footer navigation">
      
        <a href="use_pass_infra.html" class="btn btn-neutral float-right" title="How to Use TVM Pass Infra" accesskey="n" rel="next">下一个 <span class="fa fa-arrow-circle-right"></span></a>
      
      
        <a href="index.html" class="btn btn-neutral float-left" title="Extend TVM" accesskey="p" rel="prev"><span class="fa fa-arrow-circle-left"></span> 上一个</a>
      
    </div>

<div id="button" class="backtop"><img src="../../_static//img/right.svg" alt="backtop"/> </div>
<section class="footerSec">
    <div class="footerHeader">
      <ul class="d-flex align-md-items-center justify-content-between flex-column flex-md-row">
        <li class="copywrite d-flex align-items-center">
          <h5 id="copy-right-info">© 2020 Apache Software Foundation | All right reserved</h5>
        </li>
      </ul>

    </div>

    <ul>
      <li class="footernote">Copyright © 2020 The Apache Software Foundation. Apache TVM, Apache, the Apache feather, and the Apache TVM project logo are either trademarks or registered trademarks of the Apache Software Foundation.</li>
    </ul>

</section>
</footer>
        </div>
      </div>

    </section>

  </div>
  

    <script src="https://cdnjs.cloudflare.com/ajax/libs/popper.js/1.12.9/umd/popper.min.js" integrity="sha384-ApNbgh9B+Y1QKtv3Rn7W3mgPxhU9K/ScQsAP7hUibX39j7fakFPskvXusvfa0b4Q" crossorigin="anonymous"></script>
    <script src="https://maxcdn.bootstrapcdn.com/bootstrap/4.0.0/js/bootstrap.min.js" integrity="sha384-JZR6Spejh4U02d8jOt6vLEHfe/JQGiRRSQQxSfFWpi1MquVdAyjUar5+76PVCmYl" crossorigin="anonymous"></script>

  </body>
  <script type="text/javascript">
      jQuery(function () {
          SphinxRtdTheme.Navigation.enable(true);
      });
  </script>

  
  
    
    <!-- Theme Analytics -->
    <script>
    (function(i,s,o,g,r,a,m){i['GoogleAnalyticsObject']=r;i[r]=i[r]||function(){
      (i[r].q=i[r].q||[]).push(arguments)},i[r].l=1*new Date();a=s.createElement(o),
      m=s.getElementsByTagName(o)[0];a.async=1;a.src=g;m.parentNode.insertBefore(a,m)
    })(window,document,'script','https://www.google-analytics.com/analytics.js','ga');

    ga('create', 'UA-75982049-2', 'auto');
    ga('send', 'pageview');
    </script>

    
   

</body>
</html>